import argparse
import logging
import pickle
import random
import os
import json
import re
from utils1 import fix_seed, setup_data_loader, answer_cleansing, decoder_for_zetatechs, get_embedding_zetatechs, file_to_string, value_function, shapley_value
from itertools import combinations
from collections import Counter
import torch
from bert_score import score
import time
import multiprocessing
from concurrent.futures import ThreadPoolExecutor
import numpy as np
import jieba
from eval_utils import compute_bert_score, compute_safety_completeness, compute_threshold_logic, compute_logic_match

class NashDecoder:
    def __init__(self, user_information, game_goalandtips, initial_template, referee_template, 
                 answer_filter_template, model, temperature, args, CoT_template=None, 
                 tolerance=1, initial_system=None, direct_answer_trigger=None, 
                 threshold_rules=None, max_players=4, n_outer=2, n_mini=1):
        self.args = args
        self.temperature = temperature
        self.model = model
        self.direct_answer_trigger = direct_answer_trigger if direct_answer_trigger else "最终答案是"
        
        with open(user_information, 'rb') as handle:
            self.user_information = pickle.load(handle)
        self.all_player_list = list(self.user_information.keys())
        self.player_pool = {key: file_to_string(initial_template) for key in self.all_player_list}
        
        self.game_goalandtips = file_to_string(game_goalandtips)
        self.initial_system = file_to_string(initial_system) if initial_system else ""
        self.referee_template = file_to_string(referee_template)
        self.answer_filter_template = file_to_string(answer_filter_template)
        self.cot_prompt = file_to_string(CoT_template) if CoT_template else """
            任务：分析气体传感器数据并提供答案，严格按照阈值规则和输入数据进行推理。
            步骤：
            1. 提取输入数据中的关键信息（测点类型、数值、单位、报警时长、事件、历史数据）。
            2. 根据阈值规则（{threshold_rules}）判断是否超标、报警类型和报警级别。
            3. 分析事件（如放炮作业、设备标校）和历史数据，推断原因。
            4. 根据报警级别和事件，提出具体的安全建议（如撤离、通风、检查设备）。
            5. 按以下格式输出：
            是否超标：[是|否|未知]
            报警类型：[标校|超限报警|传感器断线|基站不通|未知]
            报警级别：[一级|二级|三级|未知]
            原因：[具体原因]
            建议：[具体建议]
            """
        self.threshold_rules = file_to_string(threshold_rules) if threshold_rules else "无阈值规则提供 "
        
        self.question_tag = "Q: "
        self.answer_tag = "A: "
        
        self.api_base = "https://api.zetatechs.com/v1"
        self.api_key = os.getenv("ZETA_API_KEY", "sk-p4Cown7yubjz4gETR5BGLJxBkWAQmBrCCidgt3RULy2giXFq")
        
        self.path_history = {}
        self.lambda_reward = 1.0
        self.beta = 0.1
        self.max_players = max_players
        self.tolerance = tolerance
        self.n_outer = n_outer
        self.n_mini = n_mini
        self.coalition_data = []

    def get_embedding(self, text):
        return get_embedding_zetatechs(text, self.api_key, self.api_base)

    def pre_select_players(self, query, max_players):
        print("开始预筛选玩家")
        query_emb = self.get_embedding(query)
        gas_keywords = ["甲烷", "一氧化碳", "氧气", "硫化氢", "二氧化碳", "氢气"]
        event_keywords = ["放炮作业", "设备调整", "更换传感器", "设备更换"]
        
        # CoT 提示提取任务类型和推荐角色
        cot_prompt = f"""
        任务：{query}
        请按照以下步骤分析并选择适合的角色：
        1. 提取任务中的关键信息（如气体类型、事件、历史数据）。
        2. 判断任务类型（状态识别、原因归因、策略推荐、复合任务）。
        3. 根据任务类型和关键信息，推荐适合的角色（如安全工程师、数据分析师、设备维护人员）。
        4. 分析任务层次结构，识别子任务（如状态判断、原因分析、建议生成）。
        输出格式：
        推荐角色：[角色1, 角色2, ...]
        任务类型：[状态识别|原因归因|策略推荐|复合任务]
        """
        messages = [{"role": "user", "content": cot_prompt}]
        cot_response = decoder_for_zetatechs(self.args, messages, self.args.max_length_cot // 2)
        recommended_players = re.findall(r'推荐角色：$$ ([^ $$]+)\]', cot_response)
        task_type = re.findall(r'任务类型：$$ ([^ $$]+)\]', cot_response)
        recommended_players = recommended_players[0].split(', ') if recommended_players else self.all_player_list
        task_type = task_type[0] if task_type else "复合任务"

        scores = []
        for player in self.all_player_list:
            player_desc = self.user_information[player].lower()
            score = 1.0
            
            # 关键词加权，结合任务类型
            if any(kw in player_desc for kw in gas_keywords + event_keywords):
                if task_type == "状态识别" and "数据" in player_desc:
                    score *= 1.3
                elif task_type == "原因归因" and "工程师" in player_desc:
                    score *= 1.3
                elif task_type == "策略推荐" and "安全" in player_desc:
                    score *= 1.5
                elif task_type == "复合任务" and any(kw in player_desc for kw in ["安全", "工程师", "数据"]):
                    score *= 1.2
            
            # 语义相似度计算
            if query_emb:
                player_emb = self.get_embedding(player_desc)
                if player_emb:
                    similarity = sum(a * b for a, b in zip(query_emb, player_emb)) / (
                        (sum(a * a for a in query_emb) ** 0.5) * (sum(b * b for b in player_emb) ** 0.5) + 1e-10)
                    score *= (similarity + 1.0)
            
            if player in recommended_players:
                score *= 1.3
            
            scores.append((player, score))
        
        # 语义邻接图推断新角色
        if len(scores) < max_players:
            for player in self.all_player_list:
                if player not in [p[0] for p in scores]:
                    player_desc = self.user_information[player].lower()
                    related_score = max([sum(1 for kw in extract_keywords(player_desc) if kw in extract_keywords(self.user_information[p].lower())) for p, _ in scores])
                    scores.append((player, related_score * 0.8))
        
        scores.sort(key=lambda x: x[1], reverse=True)
        selected = [p[0] for p in scores[:max_players]]
        
        if len(selected) < 2:
            remaining = [p for p in self.all_player_list if p not in selected]
            selected.extend(remaining[:2 - len(selected)])
        
        print(f"预筛选玩家: {selected}, 任务类型: {task_type}")
        logging.info(f"预筛选玩家: {selected}, 任务类型: {task_type}")
        return selected
    
    def compute_reward(self, player_id, query, answer, all_answers):
        query_lower = query.lower()
        answer_str = json.dumps(answer, ensure_ascii=False) if isinstance(answer, dict) else answer
        answer_lower = answer_str.lower()
        
        w1 = 0.7 if any(kw in query_lower for kw in ["甲烷", "一氧化碳", "氧气", "硫化氢", "二氧化碳", "氢气"]) else 0.5
        w2 = 0.3 if any(kw in query_lower for kw in ["甲烷", "一氧化碳", "氧气", "硫化氢", "二氧化碳", "氢气"]) else 0.5
        
        task_relevance = 1.0
        if all(key in answer_lower for key in ["是否超标", "报警类型", "报警级别", "原因", "建议"]):
            task_relevance = 1.5
        elif any(key in answer_lower for key in ["是否超标", "报警类型", "报警级别"]):
            task_relevance = 1.2
        else:
            task_relevance = 0.8
        
        event_keywords = ["设备调整", "放炮作业", "更换传感器", "设备更换"]
        if any(kw in answer_lower for kw in event_keywords) and any(kw in query_lower for kw in event_keywords):
            task_relevance *= 1.3
        
        safety_keywords = ["停工", "撤离", "通风", "检查", "修复", "校准", "新鲜风流", "更换接线盒"]
        if any(kw in answer_lower for kw in safety_keywords):
            task_relevance *= 1.4
        
        acc = sum(1 for a in all_answers if json.dumps(a, ensure_ascii=False).lower() == answer_lower) / len(all_answers) * task_relevance
        if query not in self.path_history:
            self.path_history[query] = {}
        if player_id not in self.path_history[query]:
            self.path_history[query][player_id] = [0.0, 0]
        
        self.path_history[query][player_id][0] = (self.path_history[query][player_id][0] * 
                                                 self.path_history[query][player_id][1] + acc) / (
                                                 self.path_history[query][player_id][1] + 1)
        self.path_history[query][player_id][1] += 1
        
        div = 0.0
        answer_emb = self.get_embedding(answer_str)
        if answer_emb:
            for other_ans in all_answers:
                other_str = json.dumps(other_ans, ensure_ascii=False) if isinstance(other_ans, dict) else other_ans
                if other_str.lower() == answer_lower:
                    continue
                other_emb = self.get_embedding(other_str)
                if other_emb:
                    div += 1 - (sum(a * b for a, b in zip(answer_emb, other_emb)) / (
                        (sum(a * a for a in answer_emb) ** 0.5) * (sum(b * b for b in other_emb) ** 0.5)))
            div /= (len(all_answers) - 1) if len(all_answers) > 1 else 1.0
        
        player_acc = self.path_history[query][player_id][0]
        other_acc = sum(self.path_history[query][p][0] for p in self.path_history[query] if p != player_id) / (
            len(self.path_history[query]) - 1) if len(self.path_history[query]) > 1 else 0.0
        reward = self.lambda_reward * (w1 * player_acc + w2 * other_acc - self.beta * div)
        return reward

    def generate_reference_answer(self, query):
        prompt = f"""
        请根据以下阈值规则和输入数据生成参考答案，严格遵循规则：
        {self.threshold_rules}
        输入：{query}
        输出格式：
        是否超标：[是|否|未知]
        报警类型：[标校|超限报警|传感器断线|基站不通|未知]
        报警级别：[一级|二级|三级|未知]
        原因：[具体原因]
        建议：[具体建议]
        """
        messages = [{"role": "user", "content": prompt}]
        response = decoder_for_zetatechs(self.args, messages, self.args.max_length_cot // 2)
        return answer_cleansing(self.args, response, self.args.dataset)

    def point_the_optimal_answer_with_cot(self, query, answer_list, active_players):
        if not answer_list:
            return {}, ""

        # 优先选择无空字段且非默认答案的选项
        default_reason = "未提供具体原因"
        default_suggestion = "建议继续监测气体浓度并检查传感器状态"
        valid_answers = [
            a for a in answer_list
            if a and isinstance(a, dict) and
            all(a.get(key) and a.get(key) != "未知" for key in ["is_exceeded", "alarm_type", "alarm_level", "reason", "suggestion"]) and
            a.get("reason") != default_reason and
            a.get("suggestion") != default_suggestion and
            not re.search(r"^(为|是：|为：)", a.get("alarm_type", "") + a.get("reason", "") + a.get("suggestion", ""))
        ]
        if not valid_answers:
            valid_answers = [
                a for a in answer_list
                if a and isinstance(a, dict) and
                a.get("reason") != default_reason and
                a.get("suggestion") != default_suggestion
            ]
        if not valid_answers:
            valid_answers = [a for a in answer_list if a and isinstance(a, dict)]
        if not valid_answers:
            return {}, ""

        rewards = []
        reference_answer = self.generate_reference_answer(query)

        for i, ans in enumerate(valid_answers):
            reward = self.compute_reward(self.optimal_player_id, query, ans, valid_answers)

            # 计算KL散度
            ans_emb = self.get_embedding(json.dumps(ans, ensure_ascii=False))
            ref_emb = self.get_embedding(json.dumps(reference_answer, ensure_ascii=False))
            if ans_emb and ref_emb:
                kl_div = sum(a * np.log((a + 1e-10) / (b + 1e-10)) for a, b in zip(ans_emb, ref_emb))
                reward -= self.beta * kl_div

            # 惩罚冗余词语
            if any(re.search(r"^(为|是：|为：)", ans.get(key, "")) for key in ["alarm_type", "reason", "suggestion"]):
                reward -= 0.5
            # 惩罚冗长答案
            if len(ans.get("reason", "")) > 50 or len(ans.get("suggestion", "")) > 100:
                reward -= 0.3

            rewards.append(reward)

        # 投票机制优化
        answer_counts = Counter([json.dumps(a, ensure_ascii=False, sort_keys=True) for a in valid_answers])
        top_answers = answer_counts.most_common(2)
        if len(top_answers) > 1 and top_answers[0][1] - top_answers[1][1] < self.tolerance:
            max_score_idx = rewards.index(max(rewards))
            selected_answer = valid_answers[max_score_idx]
        else:
            selected_answer = json.loads(top_answers[0][0])

        # 后处理：再次规范化字段
        for key in ["alarm_type", "reason", "suggestion"]:
            selected_answer[key] = re.sub(r"^(为|是：|为：)", "", selected_answer.get(key, "")).strip()

        # 生成答案选择理由
        options = "\n".join([f"{i+1}. 是否超标：{a.get('is_exceeded', '未知')}，报警类型：{a.get('alarm_type', '未知')}，报警级别：{a.get('alarm_level', '未知')}，原因：{a.get('reason', '无')}，建议：{a.get('suggestion', '无')}"
                            for i, a in enumerate(valid_answers)])
        prompt = self.answer_filter_template.format(
            query=query,
            task_type="气体传感器数据分析",
            player=active_players[self.optimal_player_id],
            instruction=self.user_information[active_players[self.optimal_player_id]],
            options=options,
            rewards="\n".join([f"答案 {i+1} 得分：{s:.3f}" for i, s in enumerate(rewards)])
        )

        message = [{"role": "user", "content": prompt}]
        response = decoder_for_zetatechs(self.args, message, self.args.max_length_cot // 2)
        logging.debug(f"答案选择响应: {response}")

        return selected_answer, response
    
    def answer_filter(self, query, args, active_players):
        print("开始多路径答案生成")
        all_answers = []

        
        def generate_player_answer(player, sub_task):
            if sub_task == "sub_task1":
                task_instruction = f"""
                请根据以下阈值规则和输入数据判断气体传感器读数是否超标，确定报警类型，并提供技术或数据视角的原因：
                {self.threshold_rules}
                特别注意事件信息（如设备调整、放炮作业、更换传感器）和历史数据对原因的影响。原因需具体描述事件或数据异常的关联。
                输入：{query}
                输出格式：
                是否超标：[是|否|未知]
                报警类型：[标校|超限报警|传感器断线|基站不通|未知]
                原因：[具体原因]
                """
            elif sub_task == "sub_task2":
                task_instruction = f"""
                请根据以下阈值规则和输入数据确定报警级别，并提供安全视角的应对建议：
                {self.threshold_rules}
                建议需结合事件信息（如放炮作业、设备调整）、气体类型和报警时长，提出具体、可操作的安全措施（如停工、撤离、增加通风）。
                输入：{query}
                输出格式：
                报警级别：[一级|二级|三级|未知]
                建议：[具体建议]
                """
            else:
                task_instruction = f"""
                请分析气体传感器数据趋势，补充原因说明，结合历史数据和事件因素（如设备调整、放炮作业、更换传感器）。
                {self.threshold_rules}
                输入：{query}
                输出格式：
                原因：[数据趋势原因]
                """

            prompt = self.player_pool[player].format(
                player_key=self.user_information[player],
                task_instruction=task_instruction,
                query=query,
                cot_prompt=self.cot_prompt.format(threshold_rules=self.threshold_rules),
                threshold_rules=self.threshold_rules,
                expected_output="""严格按照以下格式输出：
                是否超标：[是|否|未知]
                报警类型：[标校|超限报警|传感器断线|基站不通|未知]
                报警级别：[一级|二级|三级|未知]
                原因：[具体原因]
                建议：[具体建议]"""
            )

            messages = [{"role": "user", "content": prompt}]
            response = decoder_for_zetatechs(self.args, messages, args.max_length_cot)
            logging.debug(f"玩家 {player} (sub_task={sub_task}) 原始响应: {response}")
            return (player, sub_task, response)

        def generate_generic_answer():
            prompt = f"""
            分析以下气体传感器数据并提供答案，严格按照阈值规则：
            {self.threshold_rules}
            输入：{query}
            请按照以下步骤推理：
            1. 提取输入数据中的测点类型、数值、单位、报警时长、事件和历史数据。
            2. 根据阈值规则判断是否超标、报警类型和报警级别。
            3. 分析事件和历史数据，推断原因。
            4. 根据报警级别和事件，提出具体的安全建议。
            输出格式：
            是否超标：[是|否|未知]
            报警类型：[标校|超限报警|传感器断线|基站不通|未知]
            报警级别：[一级|二级|三级|未知]
            原因：[具体原因]
            建议：[具体建议]
            """
            messages = [{"role": "user", "content": prompt}]
            response = decoder_for_zetatechs(self.args, messages, args.max_length_cot)
            logging.debug(f"通用任务原始响应: {response}")
            return ("generic", "generic_task", response)
        
        for outer_iter in range(self.n_outer):
            answers = []
            for mini_iter in range(self.n_mini):
                sub_task_answers = {}
                with ThreadPoolExecutor(max_workers=len(self.optimal_coalition) + 1) as executor:
                    futures = []
                    for player in self.optimal_coalition:
                        sub_task = self.task_assignment.get(player, "sub_task1")
                        futures.append(executor.submit(generate_player_answer, player, sub_task))
                    futures.append(executor.submit(generate_generic_answer))
                    for future in futures:
                        player, sub_task, response = future.result()
                        logging.debug(f"玩家 {player} (sub_task={sub_task}) 响应: {response}")
                        cleaned = answer_cleansing(args, response, self.args.dataset, sub_task=sub_task)
                        logging.debug(f"清理后结果: {cleaned}")
                        if cleaned and any(cleaned[key] != "未知" for key in ["is_exceeded", "alarm_type", "alarm_level"]):
                            sub_task_answers[(player, sub_task)] = cleaned

                    if sub_task_answers:
                        merged_answer = {
                            "is_exceeded": "未知",
                            "alarm_type": "未知",
                            "alarm_level": "未知",
                            "reason": "",
                            "suggestion": ""
                        }
                        for (player, sub_task), ans in sub_task_answers.items():
                            if sub_task == "sub_task1":
                                if ans.get("is_exceeded") != "未知":
                                    merged_answer["is_exceeded"] = ans["is_exceeded"]
                                if ans.get("alarm_type") != "未知":
                                    merged_answer["alarm_type"] = ans["alarm_type"]
                                if ans.get("reason"):
                                    merged_answer["reason"] = ans["reason"]
                            elif sub_task == "sub_task2":
                                if ans.get("alarm_level") != "未知":
                                    merged_answer["alarm_level"] = ans["alarm_level"]
                                if ans.get("suggestion"):
                                    merged_answer["suggestion"] = ans["suggestion"]
                            elif sub_task == "sub_task3":
                                if ans.get("reason"):
                                    merged_answer["reason"] = merged_answer["reason"] + " " + ans["reason"] if merged_answer["reason"] else ans["reason"]
                            elif sub_task == "generic_task":
                                for key in ["is_exceeded", "alarm_type", "alarm_level", "reason", "suggestion"]:
                                    if ans.get(key) and ans.get(key) != "未知":
                                        merged_answer[key] = ans[key]

                        if not merged_answer["reason"]:
                            merged_answer["reason"] = "未提供具体原因"
                        if not merged_answer["suggestion"]:
                            merged_answer["suggestion"] = "建议继续监测气体浓度并检查传感器状态"

                        answers.append(merged_answer)

                all_answers.extend(answers)
        return all_answers
    
    def confine_coalition(self, query, max_length=32):
        print("使用合作博弈论选择最合适的联盟")
        active_players = self.pre_select_players(query, self.max_players)
        shapley_values = {p: shapley_value(p, active_players, query, self.user_information,
                                        self.path_history, self.get_embedding, self.api_key,
                                        self.api_base, n_samples=10) for p in active_players}

        sub_tasks = {
            "sub_task1": "状态识别（是否超标、报警类型）",
            "sub_task2": "原因归因（原因分析）",
            "sub_task3": "策略推荐（报警级别、建议生成）"
        }

        best_coalition = None
        best_task_assignment = {}
        max_value = -float('inf')
        coalition_values = {}

        # 根据任务类型动态调整权重
        task_type = "原因归因" if "原因" in query else "状态识别" if "是否超标" in query else "策略推荐"
        w1 = 0.5 if task_type == "状态识别" else 0.3  # 准确性权重
        w2 = 0.4 if task_type == "原因归因" else 0.3  # 相关性权重
        w3 = 0.3 if task_type == "策略推荐" else 0.2  # 响应效率权重

        # 记录响应时间
        response_times = {p: [] for p in active_players}

        for size in range(2, min(3, len(active_players) + 1)):
            for coalition in combinations(active_players, size):
                coalition = list(coalition)
                start_time = time.time()
                value = value_function(query, coalition, self.user_information,
                                    self.path_history, self.get_embedding,
                                    self.api_key, self.api_base, w1=w1, w2=w2, w3=w3)
                end_time = time.time()
                coalition_values[",".join(coalition)] = value

                task_assignment = {}
                has_sub_task1 = False
                has_sub_task2 = False
                has_sub_task3 = False
                for player in coalition:
                    player_desc = self.user_information[player].lower()
                    player_acc = self.path_history.get(query, {}).get(player, [0.0, 0])[0]
                    if "安全" in player_desc and not has_sub_task3 and player_acc > 0.5:
                        task_assignment[player] = "sub_task3"
                        has_sub_task3 = True
                    elif "工程师" in player_desc and not has_sub_task1 and player_acc > 0.5:
                        task_assignment[player] = "sub_task1"
                        has_sub_task1 = True
                    elif "数据" in player_desc and not has_sub_task2 and player_acc > 0.5:
                        task_assignment[player] = "sub_task2"
                        has_sub_task2 = True
                    else:
                        task_assignment[player] = "sub_task2"

                    response_times[player].append(end_time - start_time)

                if not has_sub_task3 and coalition:
                    task_assignment[coalition[0]] = "sub_task3"
                if not has_sub_task1 and len(coalition) > 1:
                    task_assignment[coalition[1]] = "sub_task1"
                if not has_sub_task2 and len(coalition) > 2:
                    task_assignment[coalition[2]] = "sub_task2"

                if value > max_value:
                    max_value = value
                    best_coalition = coalition
                    best_task_assignment = task_assignment

        if not best_coalition:
            best_coalition = active_players[:2] if len(active_players) >= 2 else active_players + [active_players[0]]
            best_task_assignment = {best_coalition[0]: "sub_task3"}
            if len(best_coalition) > 1:
                best_task_assignment[best_coalition[1]] = "sub_task1"
            max_value = value_function(query, best_coalition, self.user_information,
                                    self.path_history, self.get_embedding,
                                    self.api_key, self.api_base, w1=w1, w2=w2, w3=w3)
            coalition_values[",".join(best_coalition)] = max_value

        # 计算平均响应时间
        avg_response_times = {p: sum(times) / len(times) if times else 0.0 for p, times in response_times.items()}
        coalition_entry = {
            "query": query,
            "active_players": active_players,
            "shapley_values": shapley_values,
            "optimal_coalition": best_coalition,
            "task_assignment": best_task_assignment,
            "max_value": max_value,
            "coalition_values": coalition_values,
            "response_times": avg_response_times
        }
        self.coalition_data.append(coalition_entry)

        self.optimal_coalition = best_coalition
        self.task_assignment = best_task_assignment
        self.optimal_player_id = active_players.index(best_coalition[0]) if best_coalition else 0
        self.active_players = active_players

        print(f"最优联盟: {best_coalition}, 任务分配: {best_task_assignment}, 联盟价值: {max_value}, 响应时间: {avg_response_times}")
        logging.info(f"最优联盟: {best_coalition}, 任务分配: {best_task_assignment}, 联盟价值: {max_value}, 响应时间: {avg_response_times}")
        logging.info(f"Shapley 值: {shapley_values}")
        logging.info(f"联盟价值: {coalition_values}")
        return self.optimal_coalition

    def save_coalition_data(self, log_dir):
        os.makedirs(log_dir, exist_ok=True)
        file_path = os.path.join(log_dir, 'coalition_data.json')
        with open(file_path, 'w', encoding='utf-8') as f:
            json.dump(self.coalition_data, f, ensure_ascii=False, indent=2)
        print(f"联盟数据已保存至 {file_path}")
        logging.info(f"联盟数据已保存至 {file_path}")
        for entry in self.coalition_data:
            logging.info(f"查询: {entry['query'][:50]}... 联盟: {entry['optimal_coalition']} 价值: {entry['max_value']} Shapley值: {entry['shapley_values']}")

    def answer_clean(self, pre_answer, args):
        return answer_cleansing(args, pre_answer, self.args.dataset)

    def reach_nash(self, query, args, max_length):
        print("选择最合适的玩家并寻求 Nash 平衡")
        self.confine_coalition(query, max_length=max_length)
        all_answers = self.answer_filter(query, args, self.active_players)
        
        if not all_answers:
            print("未生成有效答案")
            logging.info("未生成有效答案")
            return [], {}
        
        best_answer, best_rationale = self.point_the_optimal_answer_with_cot(query, all_answers, self.active_players)
        
        answer_counts = Counter([json.dumps(a, ensure_ascii=False, sort_keys=True) for a in all_answers])
        final_pred_str = max(answer_counts, key=answer_counts.get)
        final_pred = json.loads(final_pred_str)
        
        print(f"最终预测: {final_pred}")
        logging.info(f"最终预测: {final_pred}")
        return [all_answers, final_pred]

def parse_arguments():
    parser = argparse.ArgumentParser(description="气体传感器数据分析测试")
    parser.add_argument('--random_seed', type=int, default=1)
    parser.add_argument('--dataset', type=str, default='coal_mine_sensors')
    parser.add_argument('--method', type=str, default='nash_cot')
    parser.add_argument('--model', type=str, default='gpt-3.5-turbo')
    parser.add_argument('--temperature', type=float, default=0.7)
    parser.add_argument('--max_length_cot', type=int, default=512)
    parser.add_argument('--max_length_direct', type=int, default=128)
    parser.add_argument('--tolerance', type=int, default=1)
    parser.add_argument('--max_players', type=int, default=4)
    parser.add_argument('--n_outer', type=int, default=2)
    parser.add_argument('--n_mini', type=int, default=1)
    parser.add_argument('--data_path', type=str, default='dataset')
    parser.add_argument('--prompt_path', type=str, default='prompts')
    parser.add_argument('--log_path', type=str, default='loggings')
    parser.add_argument('--proj_name', type=str, default='nash_coal')
    parser.add_argument('--sub_dir', type=str, default='test_exp')
    return parser.parse_args()

def normalize_text(text):
    text = text.strip()
    text = re.sub(r'\s+', ' ', text)
    text = re.sub(r'[：:]\s*', ':', text)
    text = re.sub(r'[，,。！!？?]', '', text)
    return text.lower()

def extract_keywords(text):
    keywords = set(jieba.cut(text.lower()))
    return {k for k in keywords if k in ["甲烷", "一氧化碳", "氧气", "硫化氢", "二氧化碳", "氢气", "浓度", "阈值", "放炮", "设备调整", "更换传感器", "停工", "撤离", "通风", "检查", "修复", "校准", "监测", "安全"]}

def parse_answer_string(answer_str):
    if not isinstance(answer_str, str):
        return answer_str
    result = {}
    patterns = {
        "is_exceeded": r'是否超标[:：]?\s*(是|否|未知)',
        "alarm_type": r'报警类型[:：]?\s*([^：\n]+)',
        "alarm_level": r'报警级别[:：]?\s*([^：\n]+)',
        "reason": r'原因[:：]?\s*([^：\n]+)',
        "suggestion": r'建议[:：]?\s*([^：\n]+)'
    }
    for key, pattern in patterns.items():
        match = re.search(pattern, answer_str, re.MULTILINE)
        result[key] = match.group(1).strip() if match else ""
    return result

def evaluate_fields(pred, answer, args):
    pred = json.loads(pred) if isinstance(pred, str) else pred
    answer = parse_answer_string(answer) if isinstance(answer, str) else answer

    def normalize_field(text):
        text = text.strip().lower() if text else ""
        # 语义等价映射
        equivalence = {
            "一级": ["一级报警", "一级警报"],
            "二级": ["二级报警", "二级警报"],
            "三级": ["三级报警", "三级警报"],
            "超限报警": ["超限", "超标报警"],
            "传感器断线": ["断线", "传感器故障"],
            "基站不通": ["基站故障", "通讯中断"]
        }
        for key, synonyms in equivalence.items():
            if text in synonyms:
                return key
        return text

    is_exceeded_acc = float(normalize_field(pred.get("is_exceeded", "")) == normalize_field(answer.get("is_exceeded", "")))
    alarm_type_acc = float(normalize_field(pred.get("alarm_type", "")) == normalize_field(answer.get("alarm_type", "")))
    alarm_level_acc = float(normalize_field(pred.get("alarm_level", "")) == normalize_field(answer.get("alarm_level", "")))

    reason_pred = pred.get("reason", "").strip()
    reason_answer = answer.get("reason", "").strip()
    suggestion_pred = pred.get("suggestion", "").strip()
    suggestion_answer = answer.get("suggestion", "").strip()

    reason_bert = compute_bert_score(reason_pred, reason_answer)
    suggestion_bert = compute_bert_score(suggestion_pred, suggestion_answer)

    safety_completeness = compute_safety_completeness(suggestion_pred)
    threshold_logic = compute_threshold_logic(pred, answer)
    logic_match = compute_logic_match(pred, answer)

    result = {
        "is_exceeded_acc": is_exceeded_acc,
        "alarm_type_acc": alarm_type_acc,
        "alarm_level_acc": alarm_level_acc,
        "reason_bert": reason_bert,
        "suggestion_bert": suggestion_bert,
        "safety_completeness": safety_completeness,
        "threshold_logic": threshold_logic,
        "logic_match": logic_match
    }
    return result


def main():
    try:
        print("脚本开始运行...")
        args = parse_arguments()
        print(f"参数配置: {vars(args)}")
        
        fix_seed(args.random_seed)
        
        decoder = NashDecoder(
            user_information=os.path.join(args.prompt_path, 'player_instruction.pkl'),
            game_goalandtips=os.path.join(args.prompt_path, 'game_goalandtips.txt'),
            initial_template=os.path.join(args.prompt_path, 'initial_template.txt'),
            referee_template=os.path.join(args.prompt_path, 'referee_template.txt'),
            answer_filter_template=os.path.join(args.prompt_path, 'answer_filter_template.txt'),
            model=args.model,
            temperature=args.temperature,
            args=args,
            CoT_template=os.path.join(args.prompt_path, 'cot_template.txt'),
            tolerance=args.tolerance,
            initial_system=os.path.join(args.prompt_path, 'initial_system.txt'),
            direct_answer_trigger="最终答案是",
            threshold_rules=os.path.join(args.prompt_path, 'threshold_rules.txt'),
            max_players=args.max_players,
            n_outer=args.n_outer,
            n_mini=args.n_mini
        )
        print("NashDecoder 初始化完成")
        
        data_loader = setup_data_loader(args)
        print("数据加载器初始化完成")
        
        log_dir = os.path.join(args.log_path, args.proj_name, args.sub_dir, args.method, args.dataset, str(args.random_seed))
        os.makedirs(log_dir, exist_ok=True)
        logging.basicConfig(filename=os.path.join(log_dir, 'log.txt'), level=logging.DEBUG, 
                            format='%(asctime)s - %(message)s', encoding='utf-8')
        
        total = 0
        metrics_sum = {
            "is_exceeded_acc": 0.0,
            "alarm_type_acc": 0.0,
            "alarm_level_acc": 0.0,
            "reason_bert": 0.0,
            "suggestion_bert": 0.0,
            "safety_completeness": 0.0,
            "threshold_logic": 0.0,
            "logic_match": 0.0
        }
        results = []
        answer_diversity = Counter()
        answer_stability = {}
        
        start_time = time.time()
        for batch in data_loader:
            print(f"处理批次: {batch}")
            for question, answer in zip(batch['question'], batch['answer']):
                total += 1
                query = f"{decoder.question_tag}{question}\n{decoder.answer_tag}"
                logging.debug(f"处理查询: {query}")
                print(f"处理查询 {total}: {question}")
                
                preds = []
                for _ in range(3):
                    all_answers, pred = decoder.reach_nash(query, args, max_length=args.max_length_cot)
                    logging.debug(f"当前预测: {pred}")
                    logging.debug(f"所有答案: {all_answers}")
                    preds.append(json.dumps(pred, ensure_ascii=False))
                    answer_diversity.update([json.dumps(a, ensure_ascii=False) for a in all_answers])
                
                stability = len(set(preds)) / len(preds)
                answer_stability[question] = stability
                
                metrics = evaluate_fields(pred, answer, args)
                results.append({
                    'question': question,
                    'pred': pred,
                    'answer': answer,
                    'metrics': metrics,
                    'coalition_size': len(decoder.optimal_coalition),
                    'coalition': decoder.optimal_coalition,
                    'task_assignment': decoder.task_assignment,
                    'stability': stability
                })
                logging.info(f"问题: {question}\n预测: {pred}\n正确答案: {answer}\n评估: {metrics}\n联盟: {decoder.optimal_coalition}\n任务分配: {decoder.task_assignment}\n稳定性: {stability}")
                print(f"查询 {total} 处理完成，评估指标: {metrics}")
                for k in metrics:
                    metrics_sum[k] += metrics[k]
        
        end_time = time.time()
        logging.info(f"运行时间: {end_time - start_time:.2f} 秒")
        print(f"脚本运行完成，耗时 {end_time - start_time:.2f} 秒")
        
        avg_metrics = {k: v / total if total > 0 else 0.0 for k, v in metrics_sum.items()}
        avg_metrics['diversity'] = len(answer_diversity) / total
        avg_metrics['avg_stability'] = sum(answer_stability.values()) / len(answer_stability)
        logging.info(f"平均评估指标: {avg_metrics}")
        print(f"平均评估指标: {avg_metrics}")
        
        decoder.save_coalition_data(log_dir)
        
        with open(os.path.join(log_dir, 'results.pkl'), 'wb') as f:
            pickle.dump({'results': results, 'avg_metrics': avg_metrics}, f)
        print(f"结果已保存至: {os.path.join(log_dir, 'results.pkl')}")
        
    except Exception as e:
        print(f"脚本运行出错: {str(e)}")
        logging.error(f"脚本运行出错: {str(e)}")

if __name__ == "__main__":
    main()